{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6dd53d87-0d13-4d80-b53c-f8717432bd30",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:00.864367Z",
     "iopub.status.busy": "2022-05-20T04:43:00.863732Z",
     "iopub.status.idle": "2022-05-20T04:43:01.036886Z",
     "shell.execute_reply": "2022-05-20T04:43:01.036195Z",
     "shell.execute_reply.started": "2022-05-20T04:43:00.864275Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import cv2 as cv\n",
    "import glob\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "49efb7e7-92f8-4670-8a08-42e6658bcf61",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:01.445563Z",
     "iopub.status.busy": "2022-05-20T04:43:01.444877Z",
     "iopub.status.idle": "2022-05-20T04:43:01.455149Z",
     "shell.execute_reply": "2022-05-20T04:43:01.454466Z",
     "shell.execute_reply.started": "2022-05-20T04:43:01.445535Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "a = glob.glob('horse2zebra/trainA/*')\n",
    "b = glob.glob('horse2zebra/trainB/*')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5509af1d-7500-496e-a4ac-09d52cebe8ea",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:02.508352Z",
     "iopub.status.busy": "2022-05-20T04:43:02.507599Z",
     "iopub.status.idle": "2022-05-20T04:43:07.185828Z",
     "shell.execute_reply": "2022-05-20T04:43:07.185182Z",
     "shell.execute_reply.started": "2022-05-20T04:43:02.508309Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "data1 = np.empty((len(a), 256, 256, 3))\n",
    "for i, j in enumerate(a):\n",
    "    data1[i] = cv.imread(j)\n",
    "data2 = np.empty((len(b), 256, 256, 3))\n",
    "for i, j in enumerate(b):\n",
    "    data2[i] = cv.imread(j)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "61c05b0d-e352-48cd-a9f0-d03a07ee9007",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:07.187258Z",
     "iopub.status.busy": "2022-05-20T04:43:07.187003Z",
     "iopub.status.idle": "2022-05-20T04:43:07.199249Z",
     "shell.execute_reply": "2022-05-20T04:43:07.198650Z",
     "shell.execute_reply.started": "2022-05-20T04:43:07.187237Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "image = data1[10].astype(np.uint8)\n",
    "cv.imwrite('showPic.jpg', image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3f4b27b3-1b72-4a32-8075-9d38cc3bbade",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:09.276292Z",
     "iopub.status.busy": "2022-05-20T04:43:09.275760Z",
     "iopub.status.idle": "2022-05-20T04:43:11.270522Z",
     "shell.execute_reply": "2022-05-20T04:43:11.269742Z",
     "shell.execute_reply.started": "2022-05-20T04:43:09.276266Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "data1 = (data1 - 128) / 128\n",
    "data2 = (data2 - 128) / 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "48d59c9a-5b18-447a-a01d-7d4c9c2be768",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:11.272347Z",
     "iopub.status.busy": "2022-05-20T04:43:11.271867Z",
     "iopub.status.idle": "2022-05-20T04:43:11.735121Z",
     "shell.execute_reply": "2022-05-20T04:43:11.734395Z",
     "shell.execute_reply.started": "2022-05-20T04:43:11.272321Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f46bad98-799d-40c4-bdea-3a5af4f946b3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:12.540242Z",
     "iopub.status.busy": "2022-05-20T04:43:12.539704Z",
     "iopub.status.idle": "2022-05-20T04:43:12.543948Z",
     "shell.execute_reply": "2022-05-20T04:43:12.543225Z",
     "shell.execute_reply.started": "2022-05-20T04:43:12.540210Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "Tensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9fd16542-8f47-4b09-87e2-0c0c05c35ecf",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:13.278869Z",
     "iopub.status.busy": "2022-05-20T04:43:13.278110Z",
     "iopub.status.idle": "2022-05-20T04:43:13.554935Z",
     "shell.execute_reply": "2022-05-20T04:43:13.554289Z",
     "shell.execute_reply.started": "2022-05-20T04:43:13.278829Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "dataloader1 = torch.utils.data.DataLoader(\n",
    "    Tensor(data1).permute(0, 3, 1, 2),\n",
    "    batch_size=1,\n",
    "    shuffle=True,\n",
    "    num_workers=8,\n",
    "    drop_last=True)\n",
    "\n",
    "dataloader2 = torch.utils.data.DataLoader(\n",
    "    Tensor(data2).permute(0, 3, 1, 2),\n",
    "    batch_size=1,\n",
    "    shuffle=True,\n",
    "    num_workers=8,\n",
    "    drop_last=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "709e2115-0fbe-4c8f-97bf-a6c3c8281a9c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:14.276323Z",
     "iopub.status.busy": "2022-05-20T04:43:14.275639Z",
     "iopub.status.idle": "2022-05-20T04:43:14.283073Z",
     "shell.execute_reply": "2022-05-20T04:43:14.282325Z",
     "shell.execute_reply.started": "2022-05-20T04:43:14.276284Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ResidualBlock(nn.Module):\n",
    "\n",
    "    def __init__(self, in_features):\n",
    "\n",
    "        super().__init__()\n",
    "\n",
    "        self.conv_block = nn.Sequential(\n",
    "            nn.ReflectionPad2d(1),  # +2\n",
    "            nn.Conv2d(in_features, in_features, 3),  # -2\n",
    "            nn.InstanceNorm2d(in_features),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.ReflectionPad2d(1),  # +2\n",
    "            nn.Conv2d(in_features, in_features, 3),  # -2\n",
    "            nn.InstanceNorm2d(in_features))\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x + self.conv_block(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7d736db2-d595-422a-8ce8-149b8768763a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:14.796670Z",
     "iopub.status.busy": "2022-05-20T04:43:14.796383Z",
     "iopub.status.idle": "2022-05-20T04:43:14.805707Z",
     "shell.execute_reply": "2022-05-20T04:43:14.805073Z",
     "shell.execute_reply.started": "2022-05-20T04:43:14.796645Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "\n",
    "    def __init__(self, input_nc, output_nc, n_residual_blocks=9):\n",
    "\n",
    "        super().__init__()\n",
    "\n",
    "        model = [nn.ReflectionPad2d(3), # +6\n",
    "                 nn.Conv2d(input_nc, 64, 7), # -6\n",
    "                 nn.InstanceNorm2d(64),\n",
    "                 nn.ReLU(inplace=True)]\n",
    "\n",
    "        in_features = 64\n",
    "        out_features = 64 * 2\n",
    "\n",
    "        # 下采样\n",
    "        for _ in range(2):\n",
    "            model += [nn.Conv2d(in_features,\n",
    "                                out_features,\n",
    "                                3,\n",
    "                                stride=2,\n",
    "                                padding=1),\n",
    "                      nn.InstanceNorm2d(out_features),  # nx128x128x128 nx256x64x64\n",
    "                      nn.ReLU(inplace=True)]\n",
    "            in_features = out_features\n",
    "            out_features = in_features * 2\n",
    "\n",
    "        # Residual blocks\n",
    "        for _ in range(n_residual_blocks):\n",
    "            model.append(ResidualBlock(in_features))\n",
    "\n",
    "        # 上采样\n",
    "        out_features = in_features // 2\n",
    "        for _ in range(2):\n",
    "            model += [nn.ConvTranspose2d(in_features,  # nx128x128x128\n",
    "                                         out_features,  # nx64x256x256\n",
    "                                         3,\n",
    "                                         stride=2,\n",
    "                                         padding=1,\n",
    "                                         output_padding=1),\n",
    "                      nn.InstanceNorm2d(out_features),\n",
    "                      nn.ReLU(inplace=True)]\n",
    "            in_features, out_features = out_features, out_features // 2\n",
    "\n",
    "        model += [nn.ReflectionPad2d(3),\n",
    "                  nn.Conv2d(64, output_nc, 7),\n",
    "                  nn.Tanh()]\n",
    "\n",
    "        self.model = nn.Sequential(*model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.model(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a0bc96eb-055d-4668-91ba-fe27383fe78b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:15.654747Z",
     "iopub.status.busy": "2022-05-20T04:43:15.654044Z",
     "iopub.status.idle": "2022-05-20T04:43:15.661998Z",
     "shell.execute_reply": "2022-05-20T04:43:15.661347Z",
     "shell.execute_reply.started": "2022-05-20T04:43:15.654708Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Discriminator(nn.Module):\n",
    "\n",
    "    def __init__(self, input_nc):\n",
    "\n",
    "        super().__init__()\n",
    "\n",
    "        model = [nn.Conv2d(input_nc, 64, 4, stride=2, padding=1),\n",
    "                 nn.LeakyReLU(0.2, inplace=True)]\n",
    "\n",
    "        model += [nn.Conv2d(64, 128, 4, stride=2, padding=1),\n",
    "                  nn.InstanceNorm2d(128),\n",
    "                  nn.LeakyReLU(0.2, inplace=True)]\n",
    "\n",
    "        model += [nn.Conv2d(128, 256, 4, stride=2, padding=1),\n",
    "                  nn.InstanceNorm2d(256),\n",
    "                  nn.LeakyReLU(0.2, inplace=True)]\n",
    "\n",
    "        model += [nn.Conv2d(256, 512, 4, padding=1),\n",
    "                  nn.InstanceNorm2d(512),\n",
    "                  nn.LeakyReLU(0.2, inplace=True)]\n",
    "\n",
    "        model.append(nn.Conv2d(512, 1, 4, padding=1))\n",
    "\n",
    "        self.model = nn.Sequential(*model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.model(x)\n",
    "\n",
    "        print(x.shape)\n",
    "        1/0\n",
    "        return F.avg_pool2d(x, x.size()[2:]).view(x.size()[0], -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ccb2ae8f-1e73-429b-b4ea-06c8a289fd77",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:16.515841Z",
     "iopub.status.busy": "2022-05-20T04:43:16.515126Z",
     "iopub.status.idle": "2022-05-20T04:43:16.524992Z",
     "shell.execute_reply": "2022-05-20T04:43:16.524320Z",
     "shell.execute_reply.started": "2022-05-20T04:43:16.515801Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Namespace(epoch=0, n_epochs=200, batchSize=1, dataroot='horse2zebra1/', lr=0.0002, decay_epoch=100, size=256, input_nc=3, output_nc=3, cuda=False, n_cpu=8)\n"
     ]
    }
   ],
   "source": [
    "import argparse\n",
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument('--epoch', type=int, default=0, help='starting epoch')\n",
    "parser.add_argument('--n_epochs', type=int, default=200, help='number of epochs of training')\n",
    "parser.add_argument('--batchSize', type=int, default=1, help='size of the batches')\n",
    "parser.add_argument('--dataroot', type=str, default='horse2zebra1/', help='root directory of the dataset')\n",
    "parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate')\n",
    "parser.add_argument('--decay_epoch', type=int, default=100, help='epoch to start linearly decaying the learning rate to 0')\n",
    "parser.add_argument('--size', type=int, default=256, help='size of the data crop (squared assumed)')\n",
    "parser.add_argument('--input_nc', type=int, default=3, help='number of channels of input data')\n",
    "parser.add_argument('--output_nc', type=int, default=3, help='number of channels of output data')\n",
    "parser.add_argument('--cuda', action='store_true', help='use GPU computation')\n",
    "parser.add_argument('--n_cpu', type=int, default=8, help='number of cpu threads to use during batch generation')\n",
    "opt = parser.parse_known_args()[0]\n",
    "print(opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e33bf744-20b3-428d-9edc-ef99207b66a1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:17.696083Z",
     "iopub.status.busy": "2022-05-20T04:43:17.695372Z",
     "iopub.status.idle": "2022-05-20T04:43:17.965049Z",
     "shell.execute_reply": "2022-05-20T04:43:17.964394Z",
     "shell.execute_reply.started": "2022-05-20T04:43:17.696051Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "netG_A2B = Generator(opt.input_nc, opt.output_nc)\n",
    "netG_B2A = Generator(opt.output_nc, opt.input_nc)\n",
    "netD_A = Discriminator(opt.input_nc)\n",
    "netD_B = Discriminator(opt.output_nc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f925d61f-b346-4003-92fa-fb657a09fa06",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:18.619046Z",
     "iopub.status.busy": "2022-05-20T04:43:18.618325Z",
     "iopub.status.idle": "2022-05-20T04:43:18.622845Z",
     "shell.execute_reply": "2022-05-20T04:43:18.621959Z",
     "shell.execute_reply.started": "2022-05-20T04:43:18.619017Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    netG_A2B.cuda()\n",
    "    netG_B2A.cuda()\n",
    "    netD_A.cuda()\n",
    "    netD_B.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "4e140ccd-6457-4dac-a773-f32288ac1534",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:19.866962Z",
     "iopub.status.busy": "2022-05-20T04:43:19.866400Z",
     "iopub.status.idle": "2022-05-20T04:43:20.318231Z",
     "shell.execute_reply": "2022-05-20T04:43:20.317486Z",
     "shell.execute_reply.started": "2022-05-20T04:43:19.866932Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Discriminator(\n",
       "  (model): Sequential(\n",
       "    (0): Conv2d(3, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "    (1): LeakyReLU(negative_slope=0.2, inplace=True)\n",
       "    (2): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "    (3): InstanceNorm2d(128, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "    (4): LeakyReLU(negative_slope=0.2, inplace=True)\n",
       "    (5): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "    (6): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "    (7): LeakyReLU(negative_slope=0.2, inplace=True)\n",
       "    (8): Conv2d(256, 512, kernel_size=(4, 4), stride=(1, 1), padding=(1, 1))\n",
       "    (9): InstanceNorm2d(512, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "    (10): LeakyReLU(negative_slope=0.2, inplace=True)\n",
       "    (11): Conv2d(512, 1, kernel_size=(4, 4), stride=(1, 1), padding=(1, 1))\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from utils import weights_init_normal\n",
    "netG_A2B.apply(weights_init_normal)\n",
    "netG_B2A.apply(weights_init_normal)\n",
    "netD_A.apply(weights_init_normal)\n",
    "netD_B.apply(weights_init_normal)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "73127a7c-3e3d-4ec5-9947-7a307e3d3fc8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:21.275640Z",
     "iopub.status.busy": "2022-05-20T04:43:21.275110Z",
     "iopub.status.idle": "2022-05-20T04:43:21.279589Z",
     "shell.execute_reply": "2022-05-20T04:43:21.278856Z",
     "shell.execute_reply.started": "2022-05-20T04:43:21.275613Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "criterion_GAN = torch.nn.MSELoss()\n",
    "# criterion_GAN = torch.nn.BCELoss()\n",
    "# criterion_GAN = torch.nn.CrossEntropyLoss()\n",
    "criterion_cycle = torch.nn.L1Loss()\n",
    "criterion_identity = torch.nn.L1Loss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "be5fe845-f89c-460d-890b-e37251753c16",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:24.276699Z",
     "iopub.status.busy": "2022-05-20T04:43:24.275885Z",
     "iopub.status.idle": "2022-05-20T04:43:24.284579Z",
     "shell.execute_reply": "2022-05-20T04:43:24.283635Z",
     "shell.execute_reply.started": "2022-05-20T04:43:24.276658Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import itertools\n",
    "optimizer_G = torch.optim.Adam(itertools.chain(netG_A2B.parameters(),\n",
    "                                               netG_B2A.parameters()),\n",
    "                               lr=opt.lr, betas=(0.5, 0.999))\n",
    "optimizer_D_A = torch.optim.Adam(netD_A.parameters(), lr=opt.lr, betas=(0.5, 0.999))\n",
    "optimizer_D_B = torch.optim.Adam(netD_B.parameters(), lr=opt.lr, betas=(0.5, 0.999))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2dc65b10-d4e6-4ff7-9833-8e979b99e765",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:25.555884Z",
     "iopub.status.busy": "2022-05-20T04:43:25.555179Z",
     "iopub.status.idle": "2022-05-20T04:43:25.560352Z",
     "shell.execute_reply": "2022-05-20T04:43:25.559720Z",
     "shell.execute_reply.started": "2022-05-20T04:43:25.555849Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from utils import LambdaLR\n",
    "lr_scheduler_G = torch.optim.lr_scheduler.LambdaLR(optimizer_G, lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step)\n",
    "lr_scheduler_D_A = torch.optim.lr_scheduler.LambdaLR(optimizer_D_A, lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step)\n",
    "lr_scheduler_D_B = torch.optim.lr_scheduler.LambdaLR(optimizer_D_B, lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "d0fab0bf-2883-44cf-8acc-4be94d935981",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:28.512491Z",
     "iopub.status.busy": "2022-05-20T04:43:28.512223Z",
     "iopub.status.idle": "2022-05-20T04:43:28.517057Z",
     "shell.execute_reply": "2022-05-20T04:43:28.516385Z",
     "shell.execute_reply.started": "2022-05-20T04:43:28.512465Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "input_A = Tensor(opt.batchSize, opt.input_nc, opt.size, opt.size)\n",
    "input_B = Tensor(opt.batchSize, opt.output_nc, opt.size, opt.size)\n",
    "target_real = torch.ones(opt.batchSize)\n",
    "target_fake = torch.zeros(opt.batchSize)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f6f353a3-18c8-46e8-8ab5-b91642e2b35d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:29.620194Z",
     "iopub.status.busy": "2022-05-20T04:43:29.619529Z",
     "iopub.status.idle": "2022-05-20T04:43:29.623417Z",
     "shell.execute_reply": "2022-05-20T04:43:29.622736Z",
     "shell.execute_reply.started": "2022-05-20T04:43:29.620166Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from utils import ReplayBuffer\n",
    "fake_A_buffer = ReplayBuffer()\n",
    "fake_B_buffer = ReplayBuffer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "106a005d-6ac9-4362-987a-c88a1b3378a2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-20T04:43:30.562187Z",
     "iopub.status.busy": "2022-05-20T04:43:30.561927Z",
     "iopub.status.idle": "2022-05-20T04:43:33.183560Z",
     "shell.execute_reply": "2022-05-20T04:43:33.182661Z",
     "shell.execute_reply.started": "2022-05-20T04:43:30.562150Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "0\n",
      "torch.Size([1, 1, 30, 30])\n"
     ]
    },
    {
     "ename": "ZeroDivisionError",
     "evalue": "division by zero",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mZeroDivisionError\u001b[0m                         Traceback (most recent call last)",
      "Input \u001b[0;32mIn [21]\u001b[0m, in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;66;03m# GAN loss\u001b[39;00m\n\u001b[1;32m     19\u001b[0m fake_B \u001b[38;5;241m=\u001b[39m netG_A2B(real_A)\n\u001b[0;32m---> 20\u001b[0m pred_fake \u001b[38;5;241m=\u001b[39m \u001b[43mnetD_B\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfake_B\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     21\u001b[0m loss_GAN_A2B \u001b[38;5;241m=\u001b[39m criterion_GAN(pred_fake, target_real)\n\u001b[1;32m     23\u001b[0m fake_A \u001b[38;5;241m=\u001b[39m netG_B2A(real_B)\n",
      "File \u001b[0;32m~/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py:1110\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1106\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1107\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1108\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1109\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1110\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1111\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1112\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "Input \u001b[0;32mIn [11]\u001b[0m, in \u001b[0;36mDiscriminator.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     27\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodel(x)\n\u001b[1;32m     29\u001b[0m \u001b[38;5;28mprint\u001b[39m(x\u001b[38;5;241m.\u001b[39mshape)\n\u001b[0;32m---> 30\u001b[0m \u001b[38;5;241;43m1\u001b[39;49m\u001b[38;5;241;43m/\u001b[39;49m\u001b[38;5;241;43m0\u001b[39;49m\n\u001b[1;32m     31\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m F\u001b[38;5;241m.\u001b[39mavg_pool2d(x, x\u001b[38;5;241m.\u001b[39msize()[\u001b[38;5;241m2\u001b[39m:])\u001b[38;5;241m.\u001b[39mview(x\u001b[38;5;241m.\u001b[39msize()[\u001b[38;5;241m0\u001b[39m], \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m)\n",
      "\u001b[0;31mZeroDivisionError\u001b[0m: division by zero"
     ]
    }
   ],
   "source": [
    "for epoch in range(opt.epoch, opt.n_epochs):\n",
    "    print(epoch)\n",
    "    for i, (real_A, real_B) in enumerate(zip(dataloader1, dataloader2)):\n",
    "        if not i%10:\n",
    "            print(i)\n",
    "\n",
    "        ###### Generators A2B and B2A ######\n",
    "        \n",
    "        optimizer_G.zero_grad()\n",
    "\n",
    "        same_B = netG_A2B(real_B)\n",
    "        # print(same_B.shape, real_B.shape)\n",
    "        loss_identity_B = criterion_identity(same_B, real_B)*5.0\n",
    "        # G_B2A(A) should equal A if real A is fed\n",
    "        same_A = netG_B2A(real_A)\n",
    "        loss_identity_A = criterion_identity(same_A, real_A)*5.0\n",
    "\n",
    "        # GAN loss\n",
    "        fake_B = netG_A2B(real_A)\n",
    "        pred_fake = netD_B(fake_B)\n",
    "        loss_GAN_A2B = criterion_GAN(pred_fake, target_real)\n",
    "\n",
    "        fake_A = netG_B2A(real_B)\n",
    "        pred_fake = netD_A(fake_A)\n",
    "        loss_GAN_B2A = criterion_GAN(pred_fake, target_real)\n",
    "\n",
    "        # Cycle loss\n",
    "        recovered_A = netG_B2A(fake_B)\n",
    "        loss_cycle_ABA = criterion_cycle(recovered_A, real_A)*10.0\n",
    "\n",
    "        recovered_B = netG_A2B(fake_A)\n",
    "        loss_cycle_BAB = criterion_cycle(recovered_B, real_B)*10.0\n",
    "\n",
    "        # Total loss\n",
    "        loss_G = loss_identity_A + loss_identity_B + loss_GAN_A2B + loss_GAN_B2A + loss_cycle_ABA + loss_cycle_BAB\n",
    "        loss_G.backward()\n",
    "        \n",
    "        optimizer_G.step()\n",
    "        ###################################\n",
    "\n",
    "        ###### Discriminator A ######\n",
    "        optimizer_D_A.zero_grad()\n",
    "\n",
    "        # Real loss\n",
    "        pred_real = netD_A(real_A)\n",
    "        loss_D_real = criterion_GAN(pred_real, target_real)\n",
    "\n",
    "        # Fake loss\n",
    "        fake_A = fake_A_buffer.push_and_pop(fake_A)\n",
    "        pred_fake = netD_A(fake_A.detach())\n",
    "        loss_D_fake = criterion_GAN(pred_fake, target_fake)\n",
    "\n",
    "        # Total loss\n",
    "        loss_D_A = (loss_D_real + loss_D_fake)*0.5\n",
    "        loss_D_A.backward()\n",
    "\n",
    "        optimizer_D_A.step()\n",
    "        ###################################\n",
    "\n",
    "        ###### Discriminator B ######\n",
    "        optimizer_D_B.zero_grad()\n",
    "\n",
    "        # Real loss\n",
    "        pred_real = netD_B(real_B)\n",
    "        loss_D_real = criterion_GAN(pred_real, target_real)\n",
    "        \n",
    "        # Fake loss\n",
    "        fake_B = fake_B_buffer.push_and_pop(fake_B)\n",
    "        pred_fake = netD_B(fake_B.detach())\n",
    "        loss_D_fake = criterion_GAN(pred_fake, target_fake)\n",
    "\n",
    "        # Total loss\n",
    "        loss_D_B = (loss_D_real + loss_D_fake)*0.5\n",
    "        loss_D_B.backward()\n",
    "\n",
    "        optimizer_D_B.step()\n",
    "        ###################################\n",
    "\n",
    "    \n",
    "    lr_scheduler_G.step()\n",
    "    lr_scheduler_D_A.step()\n",
    "    lr_scheduler_D_B.step()\n",
    "\n",
    "    # Save models checkpoints\n",
    "    torch.save(netG_A2B.state_dict(), 'output/netG_A2B.pth')\n",
    "    torch.save(netG_B2A.state_dict(), 'output/netG_B2A.pth')\n",
    "    torch.save(netD_A.state_dict(), 'output/netD_A.pth')\n",
    "    torch.save(netD_B.state_dict(), 'output/netD_B.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "77a0d2a3-ac95-42ef-a2e4-709705893e97",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-19T14:45:45.225181Z",
     "iopub.status.busy": "2022-05-19T14:45:45.224536Z",
     "iopub.status.idle": "2022-05-19T14:45:45.234542Z",
     "shell.execute_reply": "2022-05-19T14:45:45.234040Z",
     "shell.execute_reply.started": "2022-05-19T14:45:45.225155Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "image = (real_A*128+128)[0].permute(1,2,0).numpy().astype(np.uint8)\n",
    "cv.imwrite('showPic.jpg', image)\n",
    "image = (same_A*128+128)[0].permute(1,2,0).detach().numpy().astype(np.uint8)\n",
    "cv.imwrite('showPic1.jpg', image)\n",
    "image = (fake_B*128+128)[0].permute(1,2,0).detach().numpy().astype(np.uint8)\n",
    "cv.imwrite('showPic2.jpg', image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "884117c7-74af-4316-b880-9ee746ab32c1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-19T14:46:29.225554Z",
     "iopub.status.busy": "2022-05-19T14:46:29.224905Z",
     "iopub.status.idle": "2022-05-19T14:46:29.234582Z",
     "shell.execute_reply": "2022-05-19T14:46:29.233965Z",
     "shell.execute_reply.started": "2022-05-19T14:46:29.225519Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "image = (same_A*128+128)[0].permute(1,2,0).detach().numpy().astype(np.uint8)\n",
    "cv.imwrite('showPic1.jpg', image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "94e70b7a-47a8-457c-9694-5abc6545e5dc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-19T14:47:36.744047Z",
     "iopub.status.busy": "2022-05-19T14:47:36.743686Z",
     "iopub.status.idle": "2022-05-19T14:47:36.755030Z",
     "shell.execute_reply": "2022-05-19T14:47:36.754276Z",
     "shell.execute_reply.started": "2022-05-19T14:47:36.744019Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "image = (fake_B*128+128)[0].permute(1,2,0).detach().numpy().astype(np.uint8)\n",
    "cv.imwrite('showPic2.jpg', image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68326254-3059-45d7-838c-782e7911a7a4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
